# -*- coding: utf-8 -*-
# @File: correct.py
# @Author: Honsia
# @Time: 7月 07, 2021
# 鱼眼相机矫正
# 将鱼眼视频矫正并抽帧保存


from PIL import Image
import math
import os
import cv2
saveimg = []

def fish_eye_dis(img,im):
    "fish eye distortion"
    width_in, height_in = img.size
    im_out = Image.new("RGB",(width_in, height_in))
    radius = max(width_in, height_in)*2
    lens = radius*2/math.pi

    for i in range(width_in):
        for j in range(height_in):
            #offset to center
            x = i - width_in/2
            y = j - height_in/2
            r = math.sqrt(x*x + y*y)
            theta = math.atan(r/radius)
            if theta<0.00001:
                k = 1
            else:
                k = lens*theta/r

            src_x = x*k
            src_y = y*k
            src_x = src_x+width_in/2
            src_y = src_y+height_in/2

            pixel = im.getpixel((src_x,src_y))
            im_out.putpixel((i,j),pixel)

    return im_out

def vedio2pic(inputdir,savedir,ff):
    video_path = inputdir
    count_ = 0
    for i, file in enumerate(os.listdir(video_path)):
        v_path = os.path.join(video_path, file)
        print('正在处理：', v_path)
        video_load = cv2.VideoCapture(v_path)
        frame_count_id = 0
        while True:
            ret, img = video_load.read()
            if ret:
                frame_count_id += 1
                if frame_count_id % ff == 0:
                    count_ += 1
                    cv2.imwrite(savedir + str(count_) +'.jpg',img)
                    inputdir = savedir + str(count_) +'.jpg'
                    im = Image.open(inputdir)
                    img_out = fish_eye_dis(im,im)
                    img_out.save(savedir + "correct_"+ str(count_) +'.jpg')
                    print("save path:", savedir + "correct_"+ str(count_) +'.jpg')
                    os.remove(inputdir)
                else:
                    pass
            else:
                break

        print("finshed")

if __name__=="__main__":
    inputdir = r"C:\Users\Honsia\PycharmProjects\fisheyeCorrect\inputVedios/"  # 待处理的视频文件夹
    savedir = r"C:\Users\Honsia\PycharmProjects\fisheyeCorrect\save/"  # 视频抽帧后，图片保存的文件夹
    ff = 50  # 帧率
    v = vedio2pic(inputdir, savedir, ff)



